#!/usr/bin/env python
"""
AMQP microbenchmark.

Run as either consumer or producer.
"""

# Imports
# -------

# Standard library
import logging
import math
import optparse
import os
import signal
import struct
import sys
import time
# Third party
from netlogger import nllog
# Local
import data_pipeline as dp

# Globals
# -------

g_producer = None
g_consumer = None
g_log = None
g_workers = [ ]

# Message contents

# Classes
# -------

class AMQPError(Exception): pass

# Functions
# ---------

def handle_signals():
    for signame in 'SIGTERM', 'SIGINT', 'SIGUSR2':
        signum = getattr(signal, signame)
        signal.signal(signum, on_kill)

def on_kill(signo, frame):
    global g_consumer, g_producer
    if g_producer:
        g_producer.close()
    if g_workers:
        for w in g_workers:
            w.close()
    time.sleep(.5)
    sys.exit(0)

def get_stage_queue(num):
    """Get name of queue for a stage.
    """
    return "STAGE%d" % (num+1)

def get_result_queue():
    """Get name of qeueue for a result.
    """
    return dp.RESULT_QUEUE

def get_queue(i, d):
    """Get i-th queue in pipeline of depth d.
    """
    if i >= d:
        q = get_result_queue()
    else:
        q = get_stage_queue(i)
    return q

def run_worker(stage_num=1, worker_num=1, host='localhost'):
    global g_workers
    log = nllog.get_logger(dp.LOG_NAME)
    # Create worker
    w = dp.Worker(host=host, worker_num=worker_num)
    # Connect worker to each stage in pipeline.
    for i in xrange(stage_num):
        log.debug("stage.init", num=i)
        q_in = get_queue(i, stage_num)
        q_out = get_queue(i + 1, stage_num)
        w.add_queue_pair(q_in, q_out)
    g_workers = [ w ]
    # Run worker
    log.info("worker.start", num=worker_num)
    w.run()
    log.info("worker.end", num=worker_num, status=0)

def run_producer(msg_size=0, msg_num=None, stage_num=1,
                 host='localhost', exec_time=1, barrier=1, prob_fail=0.0):
    """Run the message producer.
    """
    global g_producer
    q_in = get_queue(0, stage_num)
    q_out = get_queue(1, stage_num)
    try:
        g_producer = dp.MasterProducer(q_in, barrier=barrier, host=host)
    except Exception, err:
        raise AMQPError("amqp connect: %s" % err)
    g_producer.send_work(count=msg_num, body_size=msg_size, exec_time=exec_time, \
            prob_fail=prob_fail)
    g_producer.close()

def run_consumer(msg_num=None, host='localhost', genstat=False, dirout='/tmp'):
    """Run the message producer.
    """
    global g_consumer
    try:
        g_consumer = dp.MasterConsumer(num_expected=msg_num, host=host,
                                        genstat=genstat, dirout=dirout)
    except Exception, err:
        raise AMQPError("amqp connect: %s" % err)
    g_consumer.run()
    g_consumer.close()
    
def main(cmdline=None):
    nllog.get_root_logger().addHandler(logging.StreamHandler())
    log = nllog.get_logger(dp.LOG_NAME)
    log.addHandler(logging.StreamHandler())
    log.setLevel(logging.WARN)
    if cmdline is None:
        cmdline = sys.argv[1:]
    handle_signals()
    op = optparse.OptionParser(usage="%prog p(roducer)|c(onsumer)|w(orker) "
                               "[options]",
                               description=__doc__)
    op.add_option("-b", "--barrier", dest="barrier", type="int", default=1,
                  help="Number of workers for producer to wait for before " +
                    "sending work to work queue (%default)")
    op.add_option("-f", "--failure", dest="failure", type="float", default=0.0,
                  help="Probability (0 <= p <= 1) of the task failing (%default)")
    op.add_option('-g', '--generate-stats', dest='genstats',
                  action='store_true', 
                  default=False, help='Generate statistic output files')
    op.add_option('-H', '--host', action='store', dest='host',
                  default="localhost",
                  metavar='HOST', help="host to conenct to (default=localhost)")
    op.add_option("-n", "--count", dest="num", type="int", default=10,
                  help="Number of messages to send (%default)")
    op.add_option('-o', '--output-dir', dest='dirout', metavar='DIR',
                  default='/tmp',
                  help='Directory to write statistics to (default=/tmp)')
    op.add_option("-s", "--stages", dest="stages", type="int", default=1,
                  help="Number pipeline stages (%default)")
    op.add_option("-t", "--exec-time", dest="exec_time", type="float",
                metavar="SEC", default=1.0,
                help="Average 'job' time (%default)")
    op.add_option("-v", "--verbose", dest="vb", action="count", default=0,
                  help="More verbose logging to standard error")
    op.add_option("-w", "--worker", dest="worker_num", type="int", default=1,
                  help="Worker number (1 .. total #workers)")
    op.add_option("-z", "--size", dest="sz", type="int", default=1024,
                  help="Total message size (%default)", metavar="BYTES")
                  
    options, args = op.parse_args(cmdline)
    if len(args) < 1:
        op.error("Must specify mode: p=producer, c=consumer, or w=worker")
    mode_str = args[0].lower()
    func = None
    if "producer".startswith(mode_str):
        func = run_producer
        kwargs = {
            'msg_num' : options.num,
            'msg_size' : options.sz,
            'stage_num' : options.stages,
            'host' : options.host,
            'exec_time' : options.exec_time,
            'barrier' : options.barrier,
            'prob_fail' : options.failure,
            }
    elif "worker".startswith(mode_str):
        func = run_worker
        kwargs = {
            'stage_num' : options.stages,
            'host' : options.host,
            'worker_num' : options.worker_num,
            }
    elif "consumer".startswith(mode_str):
        func = run_consumer
        kwargs = {
            'msg_num' : options.num,
            'host' : options.host,
            'genstat' : options.genstats,
            'dirout' : options.dirout,
            }
    if func is None:
        op.error("Unrecognized mode '%s'" % mode_str)
    # verbosity
    if options.vb > 2:
        log.setLevel(nllog.TRACE)
    elif options.vb == 2:
        log.setLevel(logging.DEBUG)
    elif options.vb == 1:
        log.setLevel(logging.INFO)
    else:
        log.setLevel(logging.WARN)
    # keyword args
    r_ = 0
    try:
        func(**kwargs)
    except AMQPError, err:
        r_ = -1
        log.error("AMQPError", msg=str(err), code=r_)
    except IOError, err:
        r_= -2
        log.error("IOError", msg=str(err), code=r_)
    return r_

if __name__ == '__main__':
    sys.exit(main())
